#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
多格式视频转换工具
"""

import os
import sys
import argparse
import subprocess
from pathlib import Path
import logging
from typing import Union, List, Optional, Dict, Any
import time
import re
import platform
import concurrent.futures
import threading
import queue
import multiprocessing

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 支持的输入格式
SUPPORTED_INPUT_FORMATS = ['mp4', 'webm', 'avi', 'mkv', 'mov', 'wmv', 'flv', 'mpg', 'mpeg', '3gp', 'm4v']
SUPPORTED_OUTPUT_FORMATS = ['mp4', 'webm', 'avi', 'mkv', 'mov']

# 推荐编码器设置
RECOMMENDED_CODECS = {
    'mp4': {'video': 'libx264', 'audio': 'aac'},
    'webm': {'video': 'libvpx-vp9', 'audio': 'libopus'},
    'avi': {'video': 'mpeg4', 'audio': 'mp3'},
    'mkv': {'video': 'libx264', 'audio': 'aac'},
    'mov': {'video': 'libx264', 'audio': 'aac'}
}

# 硬件加速编码器
HW_ACCEL_ENCODERS = {
    'nvidia': {
        'mp4': 'h264_nvenc',
        'mkv': 'h264_nvenc',
        'mov': 'h264_nvenc'
    },
    'amd': {
        'mp4': 'h264_amf',
        'mkv': 'h264_amf',
        'mov': 'h264_amf'
    },
    'intel': {
        'mp4': 'h264_qsv',
        'mkv': 'h264_qsv',
        'mov': 'h264_qsv'
    },
    'apple': {
        'mp4': 'h264_videotoolbox',
        'mkv': 'h264_videotoolbox',
        'mov': 'h264_videotoolbox'
    }
}

# 转换速度预设
SPEED_PRESETS = {
    "ultrafast": ["-preset", "ultrafast"],
    "superfast": ["-preset", "superfast"],
    "veryfast": ["-preset", "veryfast"],
    "faster": ["-preset", "faster"],
    "fast": ["-preset", "fast"],
    "medium": ["-preset", "medium"],
    "slow": ["-preset", "slow"],
    "slower": ["-preset", "slower"],
    "veryslow": ["-preset", "veryslow"]
}

# 增加速度预设描述映射，用于UI显示
SPEED_PRESET_DESCRIPTIONS = {
    "ultrafast": "极速转换，质量最低",
    "superfast": "超快速，低质量",
    "veryfast": "非常快，较低质量",
    "faster": "较快，中低质量",
    "fast": "快速，较好质量",
    "medium": "中等速度和质量，推荐",
    "slow": "慢速，高质量",
    "slower": "较慢，更高质量",
    "veryslow": "非常慢，最高质量输出"
}

class ConversionError(Exception):
    """视频转换错误类"""
    pass

def check_ffmpeg() -> bool:
    """
    检查系统是否安装了FFmpeg
    
    返回:
        bool: 如果FFmpeg已安装返回True，否则返回False
    """
    try:
        subprocess.run(
            ["ffmpeg", "-version"], 
            stdout=subprocess.PIPE, 
            stderr=subprocess.PIPE,
            check=True
        )
        return True
    except (subprocess.SubprocessError, FileNotFoundError):
        return False

def detect_hardware_acceleration() -> Dict[str, bool]:
    """
    检测系统支持的硬件加速选项
    
    返回:
        Dict[str, bool]: 可用的硬件加速选项
    """
    hw_options = {
        "nvidia": False,  # NVIDIA GPU加速
        "amd": False,     # AMD GPU加速
        "intel": False,   # Intel QSV加速
        "apple": False    # Apple VideoToolbox加速
    }
    
    try:
        # 检查FFmpeg是否已安装
        if not check_ffmpeg():
            return hw_options
            
        # 获取FFmpeg支持的编码器和设备
        result = subprocess.run(
            ["ffmpeg", "-hide_banner", "-encoders"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=False
        )
        encoders_output = result.stdout
        
        # 检查NVIDIA GPU支持
        if "h264_nvenc" in encoders_output or "hevc_nvenc" in encoders_output:
            hw_options["nvidia"] = True
            
        # 检查AMD GPU支持
        if "h264_amf" in encoders_output or "hevc_amf" in encoders_output:
            hw_options["amd"] = True
            
        # 检查Intel QSV支持
        if "h264_qsv" in encoders_output or "hevc_qsv" in encoders_output:
            hw_options["intel"] = True
            
        # 检查Apple VideoToolbox支持
        if "h264_videotoolbox" in encoders_output:
            hw_options["apple"] = True
            
        # 补充检测CUDA和其他设备
        devices_result = subprocess.run(
            ["ffmpeg", "-hide_banner", "-hwaccels"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=False
        )
        devices_output = devices_result.stdout
        
        if "cuda" in devices_output and not hw_options["nvidia"]:
            hw_options["nvidia"] = True
            
        logger.info(f"检测到的硬件加速选项: {', '.join([k for k, v in hw_options.items() if v])}")
        
    except Exception as e:
        logger.error(f"检测硬件加速选项时出错: {e}")
        
    return hw_options

def convert_video(
    input_file: Union[str, Path], 
    output_file: Optional[Union[str, Path]] = None,
    output_format: str = 'mp4',
    video_codec: str = "libx264",
    audio_codec: str = "aac",
    video_bitrate: Optional[str] = None,
    audio_bitrate: str = "192k",
    hw_accel: Optional[str] = None,
    speed_preset: str = "medium",
    progress_callback=None,
    process_ref=None,  # 添加进程引用参数
    silent=False,
    quality=23
) -> bool:
    """
    将视频文件转换为指定格式
    
    参数:
        input_file: 输入视频文件路径
        output_file: 输出视频文件路径，如果未指定则自动生成
        output_format: 输出视频格式（不含点，如'mp4'）
        video_codec: 视频编码器
        audio_codec: 音频编码器
        video_bitrate: 视频比特率
        audio_bitrate: 音频比特率
        hw_accel: 硬件加速选项 (nvidia, amd, intel, apple)
        speed_preset: 速度预设 (ultrafast, veryfast, fast, medium, slow, veryslow)
        progress_callback: 进度回调函数，接收0-1之间的浮点数
        process_ref: 用于存储进程引用的字典，方便外部终止进程
        silent: 是否静默运行（减少日志输出）
        quality: 视频质量参数 (CRF值，越小质量越高，默认23)
        
    返回:
        bool: 转换成功返回True，否则返回False
    """
    input_path = Path(input_file)
    
    # 验证输入文件
    if not input_path.exists():
        logger.error(f"输入文件不存在: {input_path}")
        return False
    
    # 如果未指定输出文件，则自动生成
    if output_file is None:
        output_format = output_format.lstrip('.')  # 确保没有前导点
        output_file = input_path.with_suffix(f'.{output_format}')
    
    output_path = Path(output_file)
    
    # 准备FFmpeg命令
    cmd = ["ffmpeg", "-y", "-i", str(input_path)]
    
    # 应用硬件加速编码器
    if hw_accel and hw_accel in HW_ACCEL_ENCODERS and output_format.lower() in HW_ACCEL_ENCODERS[hw_accel]:
        hw_video_codec = HW_ACCEL_ENCODERS[hw_accel][output_format.lower()]
        if not silent:
            logger.info(f"使用硬件加速编码器: {hw_video_codec}")
        video_codec = hw_video_codec
        
        # 添加硬件加速选项
        if hw_accel == "nvidia":
            cmd.extend(["-hwaccel", "cuda"])
        elif hw_accel == "intel":
            cmd.extend(["-hwaccel", "qsv"])
        elif hw_accel == "amd":
            cmd.extend(["-hwaccel", "amf"])
        elif hw_accel == "apple":
            cmd.extend(["-hwaccel", "videotoolbox"])
    
    # 添加视频编码参数
    cmd.extend(["-c:v", video_codec])
    
    # 如果是指定编码器且不是硬件加速，则添加速度预设
    using_cpu_codec = video_codec in ["libx264", "libx265", "libvpx", "libvpx-vp9"]
    if speed_preset in SPEED_PRESETS and using_cpu_codec:
        cmd.extend(SPEED_PRESETS[speed_preset])
        
        # 如果没有指定比特率，且使用的是x264系列编码器，使用CRF控制质量
        if not video_bitrate and any(codec in video_codec for codec in ["264", "265", "hevc"]):
            cmd.extend(["-crf", str(quality)])
    
    # 添加视频比特率
    if video_bitrate:
        cmd.extend(["-b:v", video_bitrate])
    
    # 添加音频编码参数
    cmd.extend(["-c:a", audio_codec, "-b:a", audio_bitrate])
    
    # 添加输出文件
    cmd.append(str(output_path))
    
    try:
        if not silent:
            logger.info(f"开始转换: {input_path} -> {output_path}")
            logger.debug(f"执行命令: {' '.join(cmd)}")
        
        # 获取输入文件时长
        duration = get_video_duration(input_path)
        if duration is None:
            logger.error("无法获取视频时长")
            return False
        
        # 创建进度跟踪器
        progress_tracker = ProgressTracker(duration, progress_callback)
        
        # 执行FFmpeg命令
        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE if silent else None,
            stderr=subprocess.PIPE,
            universal_newlines=True,
            encoding='utf-8',
            errors='replace'
        )
        
        # 保存进程引用
        if process_ref is not None and isinstance(process_ref, dict):
            process_ref['process'] = process
        
        # 实时读取输出并更新进度
        while True:
            line = process.stderr.readline()
            if not line and process.poll() is not None:
                break
            if line:
                progress_tracker.update(line)
        
        # 检查转换结果
        if process.returncode == 0:
            if not silent:
                logger.info(f"转换成功: {output_path}")
            return True
        elif process.returncode is None or process.returncode < 0:
            # 进程被终止
            logger.warning("转换过程被用户终止")
            return False
        else:
            logger.error(f"转换失败: {output_path}")
            return False
    except Exception as e:
        logger.error(f"转换过程出错: {str(e)}")
        return False

def get_video_duration(input_file: Union[str, Path]) -> Optional[float]:
    """
    获取视频文件的时长（秒）
    
    参数:
        input_file: 输入视频文件路径
        
    返回:
        float: 视频时长（秒），如果无法获取则返回None
    """
    try:
        cmd = [
            "ffprobe", 
            "-v", "error", 
            "-show_entries", "format=duration", 
            "-of", "default=noprint_wrappers=1:nokey=1", 
            str(input_file)
        ]
        
        result = subprocess.run(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=True
        )
        
        duration = float(result.stdout.strip())
        return duration
    except (subprocess.SubprocessError, ValueError) as e:
        logger.error(f"获取视频时长失败: {str(e)}")
        return None

class ProgressTracker:
    """进度跟踪器"""
    def __init__(self, duration, callback=None):
        self.duration = duration
        self.callback = callback
        self.start_time = time.time()
    
    def update(self, line):
        """更新进度"""
        if not self.callback:
            return
            
        # 解析FFmpeg输出中的时间信息
        time_match = re.search(r'time=(\d{2}):(\d{2}):(\d{2})\.(\d{2})', line)
        if time_match:
            hours, minutes, seconds, milliseconds = map(int, time_match.groups())
            current_time = hours * 3600 + minutes * 60 + seconds + milliseconds / 100
            
            # 计算进度百分比
            if self.duration > 0:
                progress = min(current_time / self.duration, 1.0)
                self.callback(progress)

# 为了保持向后兼容，保留原函数名称
def convert_webm_to_mp4(*args, **kwargs):
    """将WebM文件转换为MP4格式 (兼容旧版本)"""
    return convert_video(*args, **kwargs)

def _worker_process(task_args):
    """工作进程函数，用于并行批量转换时处理单个文件
    
    参数:
        task_args: 包含转换任务参数的元组
        
    返回:
        bool: 转换是否成功
    """
    try:
        input_file, output_file, output_format, video_codec, audio_codec, video_bitrate, \
        audio_bitrate, hw_accel, speed_preset = task_args
        
        # 每个工作进程独立使用convert_video函数进行转换
        return convert_video(
            input_file=input_file,
            output_file=output_file,
            output_format=output_format,
            video_codec=video_codec,
            audio_codec=audio_codec,
            video_bitrate=video_bitrate,
            audio_bitrate=audio_bitrate,
            hw_accel=hw_accel,
            speed_preset=speed_preset,
            silent=True  # 工作进程中减少输出
        )
    except Exception as e:
        logger.error(f"工作进程处理文件 {input_file} 时出错: {e}")
        return False

def batch_convert(input_dir, output_dir=None, output_format="mp4", file_extensions=None, 
                 recursive=False, video_codec="libx264", audio_codec="aac", 
                 video_bitrate=None, audio_bitrate="192k", hw_accel=None, 
                 speed_preset="medium", progress_callback=None, max_workers=None,
                 preserve_structure=True, input_formats=None):
    """批量转换视频文件

    参数:
        input_dir (str): 输入目录
        output_dir (str, optional): 输出目录。如果为None，将使用输入目录
        output_format (str, optional): 输出视频格式。默认为"mp4"
        file_extensions (list, optional): 要处理的文件扩展名列表。默认为所有支持的格式
        recursive (bool, optional): 是否递归处理子目录。默认为False
        video_codec (str, optional): 视频编码器。默认为"libx264"
        audio_codec (str, optional): 音频编码器。默认为"aac"
        video_bitrate (str, optional): 视频比特率。默认为None
        audio_bitrate (str, optional): 音频比特率。默认为"192k"
        hw_accel (str, optional): 硬件加速选项。可为"nvidia", "amd", "intel", "apple"或None
        speed_preset (str, optional): 转换速度预设。默认为"medium"
        progress_callback (callable, optional): 进度回调函数
        max_workers (int, optional): 最大工作进程数。默认为根据CPU核心数自动设置
        preserve_structure (bool, optional): 是否保留目录结构。默认为True
        input_formats (list, optional): 仅处理指定的输入格式列表。默认处理所有支持的格式
        
    返回:
        tuple: (成功转换的文件数, 总文件数)
    """
    if not os.path.isdir(input_dir):
        logger.error(f"输入目录不存在: {input_dir}")
        return 0, 0
    
    # 确定输出目录
    if output_dir is None:
        output_dir = input_dir
    
    # 创建输出目录（如果不存在）
    os.makedirs(output_dir, exist_ok=True)
    
    # 如果未指定文件扩展名，使用所有支持的格式
    if file_extensions is None:
        file_extensions = SUPPORTED_INPUT_FORMATS
    else:
        # 确保扩展名没有点号
        file_extensions = [ext.lstrip('.').lower() for ext in file_extensions]
    
    # 如果指定了输入格式过滤，则应用过滤
    if input_formats:
        input_formats = [fmt.lstrip('.').lower() for fmt in input_formats]
        file_extensions = [ext for ext in file_extensions if ext in input_formats]
    
    # 转换任务列表，每个元素为 (input_file, output_file, other_params...)
    conversion_tasks = []
    
    # 查找所有符合条件的视频文件
    input_dir_path = Path(input_dir)
    total_files = 0
    
    # 遍历文件
    if recursive:
        # 递归处理
        for root, _, files in os.walk(input_dir):
            for file in files:
                file_path = os.path.join(root, file)
                file_ext = os.path.splitext(file)[1].lstrip('.').lower()
                
                if file_ext in file_extensions:
                    # 计算输出文件路径
                    rel_path = os.path.relpath(root, input_dir) if preserve_structure else ""
                    
                    if preserve_structure:
                        out_dir = os.path.join(output_dir, rel_path)
                        os.makedirs(out_dir, exist_ok=True)
                        output_file = os.path.join(out_dir, f"{os.path.splitext(file)[0]}.{output_format}")
                    else:
                        output_file = os.path.join(output_dir, f"{os.path.splitext(file)[0]}.{output_format}")
                    
                    # 添加转换任务
                    conversion_tasks.append((
                        file_path, output_file, output_format, video_codec, audio_codec,
                        video_bitrate, audio_bitrate, hw_accel, speed_preset
                    ))
                    total_files += 1
    else:
        # 非递归处理，仅处理当前目录
        for file in os.listdir(input_dir):
            file_path = os.path.join(input_dir, file)
            
            if os.path.isfile(file_path):
                file_ext = os.path.splitext(file)[1].lstrip('.').lower()
                
                if file_ext in file_extensions:
                    output_file = os.path.join(output_dir, f"{os.path.splitext(file)[0]}.{output_format}")
                    
                    # 添加转换任务
                    conversion_tasks.append((
                        file_path, output_file, output_format, video_codec, audio_codec,
                        video_bitrate, audio_bitrate, hw_accel, speed_preset
                    ))
                    total_files += 1
    
    if total_files == 0:
        logger.warning(f"未找到符合条件的视频文件")
        return 0, 0
    
    logger.info(f"找到 {total_files} 个视频文件需要转换")
    
    # 创建进度跟踪器
    success_count = 0
    completed_count = 0
    
    # 定义进度回调函数包装器
    def update_progress():
        if progress_callback:
            progress_callback(completed_count / total_files)
    
    # 定义任务完成回调
    def task_done_callback(future):
        nonlocal success_count, completed_count
        try:
            result = future.result()
            if result:
                success_count += 1
        except Exception as e:
            logger.error(f"处理任务时出错: {e}")
        
        completed_count += 1
        update_progress()
    
    # 使用ProcessPoolExecutor进行并行处理
    import concurrent.futures
    
    # 如果未指定最大工作进程数，根据CPU核心数和文件数量自动计算
    if max_workers is None:
        import multiprocessing
        # 使用CPU核心数和文件数的较小值作为工作进程数
        max_workers = min(multiprocessing.cpu_count(), total_files)
    
    logger.info(f"使用 {max_workers} 个工作进程进行并行转换")
    
    with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
        # 提交所有任务
        future_to_file = {executor.submit(_worker_process, task): task[0] for task in conversion_tasks}
        
        # 添加回调
        for future in future_to_file:
            future.add_done_callback(task_done_callback)
        
        # 等待所有任务完成
        concurrent.futures.wait(future_to_file.keys())
    
    # 确保最终进度为100%
    if progress_callback:
        progress_callback(1.0)
    
    logger.info(f"批量转换完成。成功: {success_count}/{total_files}")
    
    return success_count, total_files

def main():
    """主程序入口"""
    # 解析命令行参数
    parser = argparse.ArgumentParser(description="多格式视频转换工具")
    
    # 子命令
    subparsers = parser.add_subparsers(dest="command", help="命令")
    
    # 单文件转换命令
    convert_parser = subparsers.add_parser("convert", help="转换单个视频文件")
    convert_parser.add_argument("input_file", help="输入视频文件路径")
    convert_parser.add_argument("-o", "--output", help="输出视频文件路径", default=None)
    convert_parser.add_argument("--output-format", help="输出视频格式", default="mp4", choices=SUPPORTED_OUTPUT_FORMATS)
    convert_parser.add_argument("--video-codec", help="视频编码器", default="libx264")
    convert_parser.add_argument("--audio-codec", help="音频编码器", default="aac")
    convert_parser.add_argument("--video-bitrate", help="视频比特率", default=None)
    convert_parser.add_argument("--audio-bitrate", help="音频比特率", default="192k")
    convert_parser.add_argument("--hw-accel", help="硬件加速选项 (nvidia, amd, intel, apple)", choices=["nvidia", "amd", "intel", "apple"], default=None)
    convert_parser.add_argument("--speed-preset", help="转换速度预设", choices=list(SPEED_PRESETS.keys()), default="medium")
    
    # 批量转换命令
    batch_parser = subparsers.add_parser("batch", help="批量转换视频文件")
    batch_parser.add_argument("input_dir", help="输入目录")
    batch_parser.add_argument("-o", "--output-dir", help="输出目录", default=None)
    batch_parser.add_argument("--output-format", help="输出视频格式", default="mp4", choices=SUPPORTED_OUTPUT_FORMATS)
    batch_parser.add_argument("-e", "--extensions", help="要处理的文件扩展名（逗号分隔）", default=None)
    batch_parser.add_argument("-r", "--recursive", help="递归处理子目录", action="store_true")
    batch_parser.add_argument("--video-codec", help="视频编码器", default="libx264")
    batch_parser.add_argument("--audio-codec", help="音频编码器", default="aac")
    batch_parser.add_argument("--video-bitrate", help="视频比特率", default=None)
    batch_parser.add_argument("--audio-bitrate", help="音频比特率", default="192k")
    batch_parser.add_argument("--hw-accel", help="硬件加速选项 (nvidia, amd, intel, apple)", choices=["nvidia", "amd", "intel", "apple"], default=None)
    batch_parser.add_argument("--speed-preset", help="转换速度预设", choices=list(SPEED_PRESETS.keys()), default="medium")
    batch_parser.add_argument("--max-workers", help="最大工作进程数", type=int, default=None)
    
    args = parser.parse_args()
    
    # 检查FFmpeg是否已安装
    if not check_ffmpeg():
        logger.error("FFmpeg未安装或不在系统路径中。请安装FFmpeg后再试。")
        sys.exit(1)
    
    # 根据子命令执行相应功能
    if args.command == "convert":
        success = convert_video(
            args.input_file,
            args.output,
            output_format=args.output_format,
            video_codec=args.video_codec,
            audio_codec=args.audio_codec,
            video_bitrate=args.video_bitrate,
            audio_bitrate=args.audio_bitrate,
            hw_accel=args.hw_accel,
            speed_preset=args.speed_preset
        )
        sys.exit(0 if success else 1)
    
    elif args.command == "batch":
        # 处理输入的文件扩展名列表
        extensions = None
        if args.extensions:
            extensions = [ext.strip() for ext in args.extensions.split(',')]
        
        batch_convert(
            args.input_dir,
            args.output_dir,
            output_format=args.output_format,
            file_extensions=extensions,
            recursive=args.recursive,
            video_codec=args.video_codec,
            audio_codec=args.audio_codec,
            video_bitrate=args.video_bitrate,
            audio_bitrate=args.audio_bitrate,
            hw_accel=args.hw_accel,
            speed_preset=args.speed_preset,
            max_workers=args.max_workers
        )
    
    else:
        parser.print_help()
        sys.exit(1)

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        logger.info("用户中断，退出程序")
        sys.exit(0)
    except Exception as e:
        logger.exception("程序运行出错")
        sys.exit(1) 